import hashlib
import json
import uuid
from typing import Union, List
import chromadb
from chromadb.utils import embedding_functions



def deterministic_uuid(content: Union[str, bytes]) -> str:
    """Creates deterministic UUID on hash value of string or byte content.

    Args:
        content: String or byte representation of data.

    Returns:
        UUID of the content.
    """
    if isinstance(content, str):
        content_bytes = content.encode("utf-8")
    elif isinstance(content, bytes):
        content_bytes = content
    else:
        raise ValueError(f"Content type {type(content)} not supported !")

    hash_object = hashlib.sha256(content_bytes)
    hash_hex = hash_object.hexdigest()
    namespace = uuid.UUID("00000000-0000-0000-0000-000000000000")
    content_uuid = str(uuid.uuid5(namespace, hash_hex))

    return content_uuid

default_ef = embedding_functions.DefaultEmbeddingFunction()
def generate_embedding(data: str, **kwargs) -> List[float]:
    embedding = default_ef([data])
    if len(embedding) == 1:
        return embedding[0]
    return embedding

chroma_client = chromadb.PersistentClient(path="./db")
collection = chroma_client.get_or_create_collection(name="my_collection")
documents="这是一个文档，大大大大大大"
id = deterministic_uuid(documents) + "-ddl"
collection.add(
    documents=documents,
    embeddings=generate_embedding(documents),
    ids=id
)

# 获取一个存在的Collection对象
#collection = chroma_client.get_collection("testname")

# 如果不存在就创建collection对象，一般用这个更多一点
#collection = chroma_client.get_or_create_collection("testname")

all_collection = collection.get()
print(f"查询文档所有数据：{all_collection}")
def get_similar_question_sql(question: str, **kwargs) -> list:
    # 准备查询条件
    where_condition = {
    }
    query_results = collection.query(query_texts=[question],
                                     where=where_condition,
                                     n_results=1)
    print(f"查询结果：{query_results}")
    if query_results is None:
        return []

    print(f"查询距离：{query_results['distances']}")
    if "documents" in query_results:
        documents = query_results["documents"]

        if len(documents) == 1 and isinstance(documents[0], list):
            try:
                documents = [json.loads(doc) for doc in documents[0]]
            except Exception as e:
                return documents[0]

        return documents


get_similar_question_sql("查询所有数据")

